import os
import json
from utils import load_model_class
from log import logger
import logging

class TrainWorker(object):
    def __init__(self, **params):
        self._model_inst = None
        self._train_dataset = None
        self._validate_dataset = None
        self._train_params = {}
        self._model_params = {}
        self._model_ckpt = None
        self._save_param_path = None
        self._save_checkpoint_path = None

    def train(self):
        logger.log(logging.INFO, 'Loading Dataset')
        self._load_dataset()
        logger.log(logging.INFO, 'Loading Model')
        self._load_model()
        logger.log(logging.INFO, 'Loading Model Params')
        self._load_model_params()
        logger.log(logging.INFO, 'Loading Train Params')
        self._load_train_params()
        logger.log(logging.INFO, 'Loading Checkpoint Path')
        self._load_model_checkpoint()
        logger.log(logging.INFO, 'Loading Save Path')
        self._load_save_path()
        logger.log(logging.INFO, 'Loading Checkpoint')
        if self._model_ckpt is not None:
            self._model_inst.load_checkpoint(self._model_ckpt)
        logger.log(logging.INFO, 'Training...')
        self._model_inst.train(self._train_dataset, **self._train_params)
        logger.log(logging.INFO, 'Evaluating Training...')
        score = self._model_inst.evaluate(self._validate_dataset)
        logger.log(logging.INFO, 'Score is '+str(score))
        self._model_inst.save(self._save_param_path)
        logger.log(logging.INFO, 'Destroy Model')
        if self._save_checkpoint_path is not None:
            self._model_inst.save_checkpoint(self._save_checkpoint_path)
        self._model_inst.destroy()

    def _load_model_params(self):
        params_str = os.environ.get('MODEL_PARAMS', '{}')
        self._model_params = json.loads(params_str)

    def _load_train_params(self):
        params_str = os.environ.get('TRAIN_PARAMS', '{}')
        self._train_params = json.loads(params_str)

    def _load_dataset(self):
        self._train_dataset = '/root/dataset/'+os.environ.get('TRAIN_DATASET', None)
        self._validate_dataset = '/root/dataset/'+os.environ.get('VALIDATE_DATASET', None)

    def _load_model(self):
        model = load_model_class('model.'+os.environ.get('MODEL_FILE', 'model'), os.environ.get('MODEL_CLASS','Model'))
        self._model_inst = model(**self._model_params)

    def _load_model_checkpoint(self):
        self._model_ckpt = None if os.environ.get('MODEL_CHECKPOINT', None) is None else '/root/checkpoint/' + os.environ.get('MODEL_CHECKPOINT', None)

    def _load_save_path(self):
        self._save_param_path = '/root/params/'+os.environ.get('MODEL_SAVE_PATH')

    def _load_save_checkpoint_path(self):
        self._save_checkpoint_path = None if os.environ.get('MODEL_SAVE_CHECKPOINT_PATH') is None else '/root/checkpoint/'+os.environ.get('MODEL_SAVE_CHECKPOINT_PATH')